import cv2
import os


rtsp_url = ("rtsp://admin:zx123456@192.168.1.63:554/h265/ch1/main/av_stream")
# 打开 RTSP 流
cap = cv2.VideoCapture(rtsp_url)


# 检查是否成功打开流
if not cap.isOpened():
    print("Error: Unable to open RTSP stream.")
    exit()

# 获取视频帧率和尺寸
fps = int(cap.get(cv2.CAP_PROP_FPS))
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_size = (frame_width, frame_height)

# 创建保存路径
output_dir = "output"
img_dir = "img"
os.makedirs(output_dir, exist_ok=True)

# 初始化视频写入对象
video_filename = os.path.join(output_dir, "output_video1.avi")
fourcc = cv2.VideoWriter_fourcc(*'XVID')
video_writer = cv2.VideoWriter(video_filename, fourcc, fps, frame_size)

print("Press 'v' to save the current frame, 'q' to quit.")

frame_count = 0  # 保存帧计数器

try:
    while True:
        ret, frame = cap.read()
        if not ret:
            print("Failed to grab frame, stopping capture.")
            break

        # 将当前帧写入视频文件
        video_writer.write(frame)

        # 显示当前帧
        cv2.imshow("RTSP Stream", frame)

        # 检测键盘输入
        key = cv2.waitKey(1) & 0xFF

        # 按下 'q' 键退出
        if key == ord('q'):
            print("Exiting...")
            break

        # 按下 'v' 键保存当前帧
        if key == ord('v'):
            frame_count += 1
            frame_filename = os.path.join(output_dir, f"frame_{frame_count:04d}.jpg")
            cv2.imwrite(frame_filename, frame)
            print(f"Saved frame to {frame_filename}")

finally:
    # 释放资源
    cap.release()
    video_writer.release()
    cv2.destroyAllWindows()
    print(f"Video saved to {video_filename}")